import graph_utils
import networkx as nx
import pylab as pl

G = nx.Graph()

G.add_edge('a', 'e')
G.add_edge('c', 'e')
G.add_edge('b', 'e')
G.add_edge('d', 'e')

nx.draw(G)

Gx, U, Y, pp_list, np_list, nn_list = graph_utils.prepare_graph_for_training_new_protocl(
                                    G, removal_perc=0.5, undersample=False, random_state=0)

print 'U:\n'
print U

print '\n\nY:\n'
print Y

print '\n\npp_list:\n'
print pp_list


print '\n\nnp_list:\n'
print np_list

print '\n\nnn_list:\n'
print nn_list

print 'donnne'

train_set, y_train, test_set, y_test = graph_utils.get_train_test_sets(G, 0.5, False, 0)

print 'train set' 
print train_set

print y_train

print 'test set' 
print test_set

print y_test

train_quads = graph_utils.get_train_quads(train_set, y_train)

print train_quads